import os
import shutil
import matplotlib.pyplot as plt


from glob import glob
from matplotlib.colors import LinearSegmentedColormap

from utils import get_xml_points, get_map_from_points, get_map_num, get_fidtm_map_num



# 测试数据、方法信息

data_name = 'DroneRGBT'

gt_path = "../datasets/DroneRGBT/Test/GT_"
rgb_path = "../datasets/DroneRGBT/Test/RGB"
tir_path = "../datasets/DroneRGBT/Test/Infrared"
gt_list = glob(os.path.join(gt_path, '*.xml'))

method_results = {
    'MC3Net': f"../other_methods/MC3Net/map_{data_name}",
    'GETANet': f"../other_methods/GETANet/map_{data_name}",
    'DEFNet': f"../other_methods/DEFNet/map_{data_name}",
    'MSDTrans': f"../other_methods/MSDTrans/map_{data_name}",
    'SONet (Ours)': f"../my_methods/fidtm_mamba_src/map_{data_name}",
}

save_path = './4_show_distribution_dronergbt'
if os.path.exists(save_path):
    shutil.rmtree(save_path)
os.makedirs(save_path)




gt_nums = []
method_nums = {}
for method in method_results.keys():
    method_nums[method] = []

for i, gt_file in enumerate(gt_list):
    print(f">>>>>{i}", end="\r")
    
    # 1.获取原始数据
    gt_name = os.path.basename(gt_file)
    et_name = gt_name.replace("R.xml", "_RGB.pt")
    rgb_file = os.path.join(rgb_path, gt_name.replace('R.xml', '.jpg'))
    tir_file = os.path.join(tir_path, gt_name.replace('R.xml', 'R.jpg'))

    rgb_img = plt.imread(rgb_file)
    tir_img = plt.imread(tir_file)
    image_shape = rgb_img.shape[:2] # (h, w)

    # 2.获取gt num
    points = get_xml_points(gt_file)
    gt_num = len(points)
    gt_nums.append(gt_num)

    # 3.获取method num
    for i, method in enumerate(method_results.keys()):
        et_file = os.path.join(method_results[method], et_name)

        # map, num
        if method == 'SONet (Ours)':
            et_map, et_num = get_fidtm_map_num(et_file)
        else:
            et_map, et_num = get_map_num(et_file, image_shape)
        
        method_nums[method].append(et_num)


fontsize = 14

plt.rcParams.update({
    'font.size': fontsize,           # 基础字体大小
    'axes.titlesize': fontsize,      # 标题字体大小
    'axes.labelsize': fontsize,      # 坐标轴标签字体大小
    'xtick.labelsize': fontsize,     # x轴刻度标签字体大小
    'ytick.labelsize': fontsize,     # y轴刻度标签字体大小
    'legend.fontsize': fontsize,     # 图例字体大小
})

for method, et_nums in method_nums.items():
    # 计算绝对误差
    absolute_error = [abs(gt_nums[i] - et_nums[i]) for i in range(len(et_nums))]

    # 创建一个自定义的渐变色
    cmap = LinearSegmentedColormap.from_list('Absolute Error', ['green', 'red'])

    # 使用自定义渐变色绘制散点图
    plt.plot([0, 200], [0, 200])
    scatter = plt.scatter(gt_nums, et_nums, c=absolute_error, cmap=cmap, alpha=0.8)

    # 设置其他属性
    plt.grid(True)
    plt.xlabel('Ground Truth Count', fontsize=fontsize, fontweight='bold')
    plt.ylabel('Estimated Count', fontsize=fontsize, fontweight='bold')

    # 颜色条字体设置
    cbar = plt.colorbar(scatter)
    cbar.set_label('Absolute Error', fontsize=fontsize)
    cbar.ax.tick_params(labelsize=12)  # 颜色条刻度字体

    plt.savefig(f'{save_path}/{method}_distribution', bbox_inches='tight', pad_inches=0.1, dpi=300)
    plt.close()




# 合并一个图

import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import numpy as np

fontsize = 14
plt.rcParams.update({
    'font.size': fontsize,           # 基础字体大小
    'axes.titlesize': fontsize,      # 标题字体大小
    'axes.labelsize': fontsize,      # 坐标轴标签字体大小
    'xtick.labelsize': fontsize,     # x轴刻度标签字体大小
    'ytick.labelsize': fontsize,     # y轴刻度标签字体大小
    'legend.fontsize': fontsize,     # 图例字体大小
})

# 创建子图布局 - 根据方法数量调整
n_methods = len(method_nums)
n_cols = n_methods  # 每行最多3列
n_rows = 1

# 设置图形大小 - 根据子图数量调整
fig_width = 6 * n_cols
fig_height = 5 * n_rows
fig, axes = plt.subplots(n_rows, n_cols, figsize=(fig_width, fig_height))

fig.suptitle('Analysis of Absolute Error Distribution in the DroneRGBT Dataset', 
             fontsize=22, fontweight='bold', y=0.95)

# 如果只有一行，确保axes是列表
if n_rows == 1:
    axes = [axes] if n_cols == 1 else axes
elif n_cols == 1:
    axes = [[ax] for ax in axes]
else:
    axes = axes.flatten()

# 计算全局的颜色范围，确保所有子图使用相同的颜色映射范围
all_absolute_errors = []
for method, et_nums in method_nums.items():
    absolute_error = [abs(gt_nums[i] - et_nums[i]) for i in range(len(et_nums))]
    all_absolute_errors.extend(absolute_error)

global_vmin = min(all_absolute_errors)
global_vmax = max(all_absolute_errors)

# 创建自定义渐变色
cmap = LinearSegmentedColormap.from_list('Absolute Error', ['green', 'yellow', 'red'])

# 绘制所有子图但不添加颜色条
scatters = []
for idx, (method, et_nums) in enumerate(method_nums.items()):
    if idx < len(axes):
        ax = axes[idx]
        absolute_error = [abs(gt_nums[i] - et_nums[i]) for i in range(len(et_nums))]
        
        ax.plot([0, 550], [0, 550], 'k--', alpha=0.7, linewidth=1)
        scatter = ax.scatter(gt_nums, et_nums, c=absolute_error, 
                           cmap=cmap, alpha=0.8, vmin=global_vmin, vmax=global_vmax, s=40)
        scatters.append(scatter)
        
        ax.grid(True, alpha=0.3)
        ax.set_xlabel('Ground Truth Count', fontsize=18)
        ax.set_ylabel('Estimated Count', fontsize=18)
        ax.set_title(f'{method}', fontsize=18, pad=10)
        ax.set_xlim([0, max(gt_nums) * 1.05])
        ax.set_ylim([0, max(et_nums) * 1.05])

# 添加一个共享的颜色条
fig.subplots_adjust(right=0.9)
cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])  # [left, bottom, width, height]
cbar = fig.colorbar(scatters[0], cax=cbar_ax)
cbar.set_label('Absolute Error', fontsize=18)
cbar.ax.tick_params(labelsize=fontsize)

plt.tight_layout(rect=[0, 0, 0.9, 1])  # 为颜色条留出空间
plt.savefig(f'{save_path}/all_methods_comparison_shared_cbar.png', 
           bbox_inches='tight', pad_inches=0.1, dpi=300)
plt.close()



